#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
AI预测测试类
按照ToDo.md要求：
1. 输入起始期数，默认往后取30（可输入）期数据
2. 从表lottery_results取数据
3. 从起始期数开始，依次遍历至2025090期
4. 按照当前规则，将获取的历史中奖数据和它们之间数据互斥关系传给AI
5. 解析AI返回信息，得出是否中奖信息
6. 记录prompt、AI分析至表ai_predictions
"""

import pymysql
import json
import requests
import time
from datetime import datetime
from typing import List, Dict, Any, Optional
from openai import OpenAI


class AIPredictionTestClass:
    """AI预测测试类"""
    
    def __init__(self, db_config: Dict[str, Any]):
        """
        初始化测试类
        
        Args:
            db_config: 数据库配置
        """
        self.db_config = db_config
        self.connection = None
        self.cursor = None
        
        # AI配置
        self.ai_config = {
            'api_url': 'https://api.deepseek.com/v1/chat/completions',
            'api_key': 'sk-d8b48b0fcd5e4ee7823d7ce71db41fde',
            'model': 'deepseek-chat'
        }
        
        # 测试配置
        self.default_periods = 30  # 默认取30期数据
        self.max_draw = '2025090'  # 最大期数
        
    def connect_database(self):
        """连接数据库"""
        try:
            self.connection = pymysql.connect(**self.db_config)
            self.cursor = self.connection.cursor(pymysql.cursors.DictCursor)
            print("✓ 数据库连接成功")
        except Exception as e:
            print(f"✗ 数据库连接失败: {e}")
            raise
    
    def close_database(self):
        """关闭数据库连接"""
        if self.cursor:
            self.cursor.close()
        if self.connection:
            self.connection.close()
        print("✓ 数据库连接已关闭")
    
    def get_lottery_data(self, start_draw: str, period_count: int = None) -> List[Dict[str, Any]]:
        """
        从lottery_results表获取数据
        
        Args:
            start_draw: 起始期数
            period_count: 期数数量，默认使用self.default_periods
            
        Returns:
            彩票数据列表
        """
        if period_count is None:
            period_count = self.default_periods
            
        try:
            sql = """
                SELECT draw_number, draw_date, 
                       red_ball_1, red_ball_2, red_ball_3, red_ball_4, red_ball_5, red_ball_6,
                       blue_ball, prize_info
                FROM lottery_results 
                WHERE draw_number >= %s 
                ORDER BY draw_date ASC 
                LIMIT %s
            """
            
            self.cursor.execute(sql, (start_draw, period_count))
            results = self.cursor.fetchall()
            
            # 格式化数据
            formatted_data = []
            for row in results:
                red_balls = [
                    row['red_ball_1'], row['red_ball_2'], row['red_ball_3'],
                    row['red_ball_4'], row['red_ball_5'], row['red_ball_6']
                ]
                
                formatted_data.append({
                    'draw_number': row['draw_number'],
                    'draw_date': row['draw_date'].isoformat() if row['draw_date'] else None,
                    'red_balls': red_balls,
                    'blue_ball': row['blue_ball'],
                    'prize_info': row['prize_info']
                })
            
            print(f"✓ 获取到 {len(formatted_data)} 期数据，从 {start_draw} 开始")
            return formatted_data
            
        except Exception as e:
            print(f"✗ 获取彩票数据失败: {e}")
            return []
    
    def analyze_mutual_exclusion(self, lottery_data: List[Dict[str, Any]]) -> Dict[str, Any]:
        """
        分析红球互斥关系
        
        Args:
            lottery_data: 彩票数据列表
            
        Returns:
            互斥关系分析结果
        """
        try:
            # 统计红球出现频率
            red_ball_freq = {}
            for i in range(1, 34):
                red_ball_freq[i] = 0
            
            # 统计红球对出现频率
            red_ball_pairs = {}
            for i in range(1, 34):
                for j in range(i + 1, 34):
                    red_ball_pairs[(i, j)] = 0
            
            # 分析历史数据
            for draw in lottery_data:
                red_balls = draw['red_balls']
                
                # 统计单个红球频率
                for ball in red_balls:
                    if ball in red_ball_freq:
                        red_ball_freq[ball] += 1
                
                # 统计红球对频率
                for i in range(len(red_balls)):
                    for j in range(i + 1, len(red_balls)):
                        ball1, ball2 = min(red_balls[i], red_balls[j]), max(red_balls[i], red_balls[j])
                        if (ball1, ball2) in red_ball_pairs:
                            red_ball_pairs[(ball1, ball2)] += 1
            
            # 找出强互斥关系（从未一起出现）
            strong_exclusions = []
            for pair, freq in red_ball_pairs.items():
                if freq == 0:
                    strong_exclusions.append({
                        'ball1': pair[0],
                        'ball2': pair[1],
                        'frequency': 0
                    })
            
            # 找出弱互斥关系（很少一起出现，频率低于阈值）
            # weak_exclusions = []
            # threshold = len(lottery_data) * 0.1  # 10%阈值
            # for pair, freq in red_ball_pairs.items():
            #     if 0 < freq <= threshold:
            #         exclusion_strength = 1 - (freq / len(lottery_data))
            #         weak_exclusions.append({
            #             'ball1': pair[0],
            #             'ball2': pair[1],
            #             'frequency': freq,
            #             'exclusion_strength': exclusion_strength
            #         })
            
            # # 按互斥强度排序
            # weak_exclusions.sort(key=lambda x: x['exclusion_strength'], reverse=True)
            
            result = {
                'total_draws': len(lottery_data),
                'strong_exclusions': strong_exclusions,
                # 'weak_exclusions': weak_exclusions,
                'red_ball_frequency': red_ball_freq
            }
            
            print(f"✓ 互斥关系分析完成: 强互斥 {len(strong_exclusions)} 个")
            return result
            
        except Exception as e:
            print(f"✗ 互斥关系分析失败: {e}")
            return {'strong_exclusions': [], 'weak_exclusions': [], 'total_draws': 0}


    def call_tongyi(self,prompt: str) -> Optional[str]:
        try:
            client = OpenAI(
                # 若没有配置环境变量，请用百炼API Key将下行替换为：api_key="sk-xxx",
                # api_key=os.getenv("sk-f025138adc8e4564824859d82e7f075c"),
                api_key="sk-f025138adc8e4564824859d82e7f075c",
                base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
            )

            completion = client.chat.completions.create(
                # 模型列表：https://help.aliyun.com/zh/model-studio/getting-started/models
                model="qwen-plus",
                messages=[
                    {"role": "system", "content": "You are a helpful assistant."},
                    {"role": "user", "content": prompt},
                ],
                # Qwen3模型通过enable_thinking参数控制思考过程（开源版默认True，商业版默认False）
                # 使用Qwen3开源版模型时，若未启用流式输出，请将下行取消注释，否则会报错
                # extra_body={"enable_thinking": False},
            )
            content = completion.choices[0].message.content
            return content

        except Exception as e:
            print(f"✗ AI API调用异常: {e}")
            return None
    def call_ai_api(self, prompt: str) -> Optional[str]:
        """
        调用AI API

        Args:
            prompt: 发送给AI的提示词

        Returns:
            AI返回的内容
        """
        try:
            headers = {
                'Content-Type': 'application/json',
                'Authorization': f'Bearer {self.ai_config["api_key"]}'
            }

            data = {
                'model': self.ai_config['model'],
                'messages': [
                    {
                        'role': 'user',
                        'content': prompt
                    }
                ],
                'temperature': 0.7,
                'max_tokens': 2000
            }

            response = requests.post(
                self.ai_config['api_url'],
                headers=headers,
                json=data,
                timeout=90
            )

            if response.status_code == 200:
                result = response.json()
                content = result['choices'][0]['message']['content']
                print("✓ AI API调用成功")
                return content
            else:
                print(f"✗ AI API调用失败: {response.status_code} - {response.text}")
                return None

        except Exception as e:
            print(f"✗ AI API调用异常: {e}")
            return None
    
    def build_ai_prompt(self, lottery_data: List[Dict[str, Any]], 
                       exclusion_analysis: Dict[str, Any],
                       custom_prompt: str = "") -> str:
        """
        构建AI提示词
        
        Args:
            lottery_data: 彩票数据
            exclusion_analysis: 互斥关系分析
            custom_prompt: 自定义提示词
            
        Returns:
            完整的AI提示词
        """
        # 构建历史数据字符串
        history_lines = []
        for draw in lottery_data:
            red_balls_str = ' '.join([f"{ball:02d}" for ball in draw['red_balls']])
            blue_ball_str = f"{draw['blue_ball']:02d}"
            history_lines.append(f"{draw['draw_number']}: {draw['draw_date']} - 红球: {red_balls_str} 蓝球: {blue_ball_str}")
        
        history_data_str = '\n'.join(history_lines)
        
        # 构建互斥关系信息
        exclusion_info = ""
        strong_exclusions = exclusion_analysis.get('strong_exclusions', [])
        # weak_exclusions = exclusion_analysis.get('weak_exclusions', [])
        
        # if strong_exclusions or weak_exclusions:
        if strong_exclusions :
            exclusion_info = "\n⚠️ 重要：互斥关系分析（必须严格遵守）：\n"
            
            if strong_exclusions:
                exclusion_info += "🚫 强互斥关系（绝对禁止同时选择）：\n"
                for i, exclusion in enumerate(strong_exclusions, 1):  # 显示全部强互斥关系
                    exclusion_info += f"  {i}. 红球{exclusion['ball1']:02d} 与 红球{exclusion['ball2']:02d}（从未一起出现）\n"
            
            # if weak_exclusions:
            #     exclusion_info += "⚠️ 弱互斥关系（建议避免同时选择）：\n"
            #     for i, exclusion in enumerate(weak_exclusions[:20], 1):  # 显示前20个
            #         exclusion_info += f"  {i}. 红球{exclusion['ball1']:02d} 与 红球{exclusion['ball2']:02d}（互斥强度：{exclusion['exclusion_strength']:.3f}）\n"
        
        # 构建完整提示词
        base_prompt = f"""请基于以下双色球历史数据分析，对下一期的双色球中奖进行预测，生成5组预测号码：

历史中奖数据分析：
- 中奖数据信息：
{history_data_str}

{exclusion_info}

"""

        if custom_prompt:
            base_prompt = custom_prompt + "\n\n" + base_prompt
            
        return base_prompt
    
    def parse_ai_response(self, response: str) -> Dict[str, Any]:
        """
        解析AI响应
        
        Args:
            response: AI返回的内容
            
        Returns:
            解析结果
        """
        try:
            predictions = []
            analysis = ""
            
            lines = response.split('\n')
            in_predictions = False
            
            for line in lines:
                line = line.strip()
                
                if '预测号码：' in line:
                    in_predictions = True
                    continue
                elif '分析说明：' in line:
                    in_predictions = False
                    continue
                
                if in_predictions and '第' in line and ('组:' in line or '组：' in line):
                    # 解析预测号码
                    import re
                    
                    # 解析红球
                    red_match = re.search(r'红球\[([^\]]+)\]', line)
                    if red_match:
                        red_balls = [int(x.strip()) for x in red_match.group(1).split(',') if x.strip().isdigit()]
                        
                        # 解析蓝球
                        blue_match = re.search(r'蓝球(\d+)', line)
                        if blue_match:
                            blue_ball = int(blue_match.group(1))
                            
                            if len(red_balls) == 6 and 1 <= blue_ball <= 16:
                                predictions.append({
                                    'red_balls': sorted(red_balls),
                                    'blue_ball': blue_ball
                                })
                
                elif not in_predictions and line:
                    analysis += line + '\n'
            
            return {
                'predictions': predictions,
                'analysis': analysis.strip(),
                'raw_response': response
            }
            
        except Exception as e:
            print(f"✗ 解析AI响应失败: {e}")
            return {
                'predictions': [],
                'analysis': '解析失败',
                'raw_response': response
            }
    
    def save_prediction_record(self, start_draw: str, end_draw: str, 
                             lottery_data: List[Dict[str, Any]],
                             ai_result: Dict[str, Any],
                             prompt: str,
                             exclusion_analysis: Dict[str, Any]) -> bool:
        """
        保存预测记录到ai_predictions表
        
        Args:
            start_draw: 开始期数
            end_draw: 结束期数
            lottery_data: 彩票数据
            ai_result: AI分析结果
            prompt: 发送给AI的prompt
            exclusion_analysis: 互斥关系分析
            
        Returns:
            是否保存成功
        """
        try:
            sql = """
                INSERT INTO ai_predictions (
                    start_draw, end_draw, data_count, predicted_numbers, 
                    analysis_result, prompt, ai_model, ai_response, 
                    analysis_process, exclusion_analysis, created_at
                ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
            """
            
            values = (
                start_draw,
                end_draw,
                len(lottery_data),
                json.dumps(ai_result['predictions'], ensure_ascii=False),
                ai_result['analysis'],
                prompt,
                'deepseek-chat',
                ai_result['raw_response'],
                f"基于{len(lottery_data)}期历史数据分析，包含互斥关系分析",
                json.dumps(exclusion_analysis, ensure_ascii=False),
                datetime.now()
            )
            
            self.cursor.execute(sql, values)
            self.connection.commit()
            
            prediction_id = self.cursor.lastrowid
            print(f"✓ 预测记录已保存，ID: {prediction_id}")
            return True
            
        except Exception as e:
            print(f"✗ 保存预测记录失败: {e}")
            self.connection.rollback()
            return False
    
    def run_prediction_test(self, start_draw: str, period_count: int = None, 
                          custom_prompt: str = "", max_iterations: int = 10):
        """
        运行预测测试
        
        Args:
            start_draw: 起始期数
            period_count: 期数数量
            custom_prompt: 自定义提示词
            max_iterations: 最大迭代次数
        """
        try:
            print(f"=== 开始AI预测测试 ===")
            print(f"起始期数: {start_draw}")
            print(f"期数数量: {period_count or self.default_periods}")
            print(f"最大期数: {self.max_draw}")
            
            current_draw = start_draw
            iteration = 0
            
            while current_draw <= self.max_draw and iteration < max_iterations:
                print(f"\n--- 第 {iteration + 1} 次迭代 ---")
                print(f"当前期数: {current_draw}")
                
                # 1. 获取彩票数据
                lottery_data = self.get_lottery_data(current_draw, period_count)
                if not lottery_data:
                    print(f"⚠ 没有找到从 {current_draw} 开始的数据，跳过")
                    # 尝试下一个期数
                    current_draw = self.get_next_draw_number(current_draw)
                    iteration += 1
                    continue
                
                # 2. 分析互斥关系
                exclusion_analysis = self.analyze_mutual_exclusion(lottery_data)
                
                # 3. 构建AI提示词
                prompt = self.build_ai_prompt(lottery_data, exclusion_analysis, custom_prompt)
                
                # 4. 调用AI API
                # ai_response = self.call_ai_api(prompt)
                # 打印开始调用时间
                print("通义模型开始调用时间:", datetime.now().strftime("%Y-%m-%d %H:%M:%S"))

                ai_response = self.call_tongyi(prompt)

                print("通义模型完成调用时间:", datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
                if not ai_response:
                    print(f"⚠ AI API调用失败，跳过当前期数")
                    current_draw = self.get_next_draw_number(current_draw)
                    iteration += 1
                    continue
                
                # 5. 解析AI响应
                ai_result = self.parse_ai_response(ai_response)
                
                # 6. 保存预测记录
                end_draw = lottery_data[-1]['draw_number'] if lottery_data else current_draw
                success = self.save_prediction_record(
                    current_draw, end_draw, lottery_data, ai_result, 
                    prompt, exclusion_analysis
                )
                
                if success:
                    print(f"✓ 第 {iteration + 1} 次迭代完成")
                    print(f"  预测数量: {len(ai_result['predictions'])}")
                    print(f"  分析长度: {len(ai_result['analysis'])}")
                else:
                    print(f"✗ 第 {iteration + 1} 次迭代失败")
                
                # 移动到下一个起始期数
                current_draw = self.get_next_draw_number(current_draw)
                iteration += 1
                
                # 添加延迟避免API限制
                time.sleep(2)
            
            print(f"\n=== 测试完成 ===")
            print(f"总迭代次数: {iteration}")
            
        except Exception as e:
            print(f"✗ 测试运行失败: {e}")
    
    def get_next_draw_number(self, current_draw: str) -> str:
        """
        获取下一个期数
        
        Args:
            current_draw: 当前期数
            
        Returns:
            下一个期数
        """
        try:
            # 假设期数格式为 YYYYNNN，其中YYYY是年份，NNN是期数
            year = current_draw[:4]
            draw_num = int(current_draw[4:])
            next_draw_num = draw_num + 1
            
            # 如果超过150期，进入下一年
            if next_draw_num > 150:
                year = str(int(year) + 1)
                next_draw_num = 1
            
            return f"{year}{next_draw_num:03d}"
            
        except Exception as e:
            print(f"✗ 计算下一个期数失败: {e}")
            return current_draw

def main():
    """主函数"""
    # 数据库配置
    db_config = {
        'host': 'localhost',
        'user': 'root',
        'password': 'mario',
        'database': 'ai_ssq',
        'charset': 'utf8mb4'
    }
    
    # 创建测试类实例
    test_class = AIPredictionTestClass(db_config)
    
    try:
        # 连接数据库
        test_class.connect_database()
        
        # 运行测试
        start_draw = '2025015'  # 起始期数
        period_count = 50       # 期数数量
        max_iterations = 10  # 最大迭代次数
        custom_prompt = f"""
分析红球的杀号规律；
- 杀号规则1：杀冷号，例如红球05大于15期未出，可对其进行杀号（并分析近n期该规则杀号成功率）；
- 杀号规则2：06近5期(指的是距离预测期数的近5期)出现了3次，属于高频号码，可将其进行杀号（并分析近n期该规则杀号成功率）；
- 杀号规则3：以上期号码作为参考，杀号（上期红球(n+1)-上期红球n）（并分析近n期该规则杀号成功率）；
- 杀号规则4: 以上期号码作为参考，杀号（上期红球分别减去蓝球的绝对值作为下期红球的杀号）（并分析近n期该规则杀号成功率）；
蓝球杀号规律：
- 蓝球杀号规则1：近5期(指的是距离预测期数的近5期)出现大于2次 的蓝球; 考虑下期杀该号码（并分析近n期该规则杀号成功率）；
- 蓝球杀号规则2：大于10期(指的是距离预测期数的近5期)以上的号码可进行杀号；（如近期大于10期的号码成功率大于70%，可执行该杀号规则）
红球选号规则：
- 红球选号规则1：优先选择5期(指的是距离预测期数的近5期)内未被杀掉的号码；
- 红球选号规则2：优先与上期没有重复的号码；
- 红球选号规则3：可优先组合连号的号码；
蓝球选号规则：
- 蓝球选号规则1：优先选择5期(指的是距离预测期数的近5期)内未被杀掉的号码；
- 蓝球选号规则2：优先与上期没有重复的号码；

🚨 互斥关系检查要求：
1. 生成预测号码前，必须仔细检查互斥关系
2. 绝对不能在同一个预测组中同时选择强互斥的号码对
3. 尽量避免在同一个预测组中同时选择弱互斥的号码对
4. 每生成一组预测后，都要验证是否违反互斥规则

请严格按照以下格式返回结果：
预测号码：
第1组:红球[1,2,3,4,5,6],蓝球04
第2组:红球[7,8,9,10,11,12],蓝球08
第3组:红球[13,14,15,16,17,18],蓝球10
第4组:红球[19,20,21,22,23,24],蓝球12
第5组:红球[25,26,27,28,29,30],蓝球16

分析说明：
[详细的分析说明(先分析红球杀号，再分析蓝球杀号)]
"""  # 自定义提示词

        
        test_class.run_prediction_test(
            start_draw=start_draw,
            period_count=period_count,
            custom_prompt=custom_prompt,
            max_iterations=max_iterations
        )
        
    except Exception as e:
        print(f"✗ 主程序执行失败: {e}")
    finally:
        # 关闭数据库连接
        test_class.close_database()

if __name__ == '__main__':
    main()
